import logging
from load_data import load_roidb
import pdb

def load_coco_test_roidb_eval(config, annotation_path, getCls2Name=False):
    """
    modify from smart
    """
    roidb   = load_roidb(roidb_path_list=config.dataset.test_roidb_path_list,
                        imglst_path_list=config.dataset.test_imglst_path_list,
                        percent         = config.dataset.test_percent,
                        filter_strategy = config.TEST.filter_strategy)
    logging.info('total num images for test: {}'.format(len(roidb)))
    
    from coco_eval import COCOEval
    imdb    = COCOEval(annotation_path)
    if 'voc' in annotation_path or 'VOC' in annotation_path \
        or 'dota' in annotation_path or 'DOTA' in annotation_path \
        or 'kitti' in annotation_path or 'KITTI' in annotation_path:
        logging.info('For {}, need sort the imageset_index'.format(annotation_path))
        print('For {}, need sort the imageset_index'.format(annotation_path))
        imdb.imageset_index     = sorted(imdb.imageset_index)
    else:
        logging.info('For {}, not change the imageset_index'.format(annotation_path))
        print('For {}, not change the imageset_index'.format(annotation_path))
    cls2name    = dict(zip(range(len(imdb.classes)), imdb.classes))

    # sample roidb
    roidb       = imdb.sample_on_imdb(roidb, filter_strategy=config.TEST.filter_strategy)
    logging.info('total num images for test after sampling: {}'.format(len(roidb)))
    
    def eval_func(**kwargs):
        # TODO
        raise NotImplementedError()
        pass
    if getCls2Name:
        return roidb, eval_func
    else:
        return roidb, eval_func, cls2name

def load_coco_train_roidb_eval(config, annotation_path):
    """
    modify from smart
    """
    roidb   = load_roidb(roidb_path_list=config.dataset.train_roidb_path_list,
                        imglst_path_list=config.dataset.train_imglst_path_list,
                        percent         = config.dataset.train_percent,
                        filter_strategy = config.TRAIN.filter_strategy)
    logging.info('total num images for train: {}'.format(len(roidb)))
    
    from coco_eval import COCOEval
    imdb    = COCOEval(annotation_path)
    if 'voc' in annotation_path or 'VOC' in annotation_path \
        or 'dota' in annotation_path or 'DOTA' in annotation_path \
        or 'kitti' in annotation_path or 'KITTI' in annotation_path:
        logging.info('For {}, need sort the imageset_index'.format(annotation_path))
        print('For {}, need sort the imageset_index'.format(annotation_path))
        imdb.imageset_index     = sorted(imdb.imageset_index)
    else:
        logging.info('For {}, not change the imageset_index'.format(annotation_path))
        print('For {}, not change the imageset_index'.format(annotation_path))
    cls2name    = dict(zip(range(len(imdb.classes)), imdb.classes))

    # sample roidb
    roidb       = imdb.sample_on_imdb(roidb, filter_strategy=config.TRAIN.filter_strategy)
    logging.info('total num images for train after sampling: {}'.format(len(roidb)))
    
    def eval_func(**kwargs):
        # TODO
        raise NotImplementedError()
        pass
    return roidb, eval_func
